// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/common/base/exceptions.h>
#include <pollux/exec/window_function.h>
#include <pollux/expression/function_signature.h>
#include <pollux/vector/flat_vector.h>

namespace kumo::pollux::window::prestosql {
    namespace {
        enum class ValueType {
            kFirst,
            kLast,
        };

        template<ValueType TValue>
        class FirstLastValueFunction : public exec::WindowFunction {
        public:
            explicit FirstLastValueFunction(
                const std::vector<exec::WindowFunctionArg> &args,
                const TypePtr &resultType,
                bool ignoreNulls,
                pollux::memory::MemoryPool *pool)
                : WindowFunction(resultType, pool, nullptr), ignoreNulls_(ignoreNulls) {
                POLLUX_CHECK_NULL(args[0].constantValue);
                valueIndex_ = args[0].index.value();

                nulls_ = allocate_nulls(0, pool_);
            }

            void resetPartition(const exec::WindowPartition *partition) override {
                partition_ = partition;
            }

            void apply(
                const BufferPtr & /*peerGroupStarts*/,
                const BufferPtr & /*peerGroupEnds*/,
                const BufferPtr &frameStarts,
                const BufferPtr &frameEnds,
                const SelectivityVector &validRows,
                int32_t resultOffset,
                const VectorPtr &result) override {
                auto numRows = frameStarts->size() / sizeof(vector_size_t);
                rowNumbers_.resize(numRows);
                if (validRows.hasSelections()) {
                    if (ignoreNulls_) {
                        setRowNumbersIgnoreNulls(validRows, frameStarts, frameEnds);
                    } else {
                        setRowNumbersRespectNulls(validRows, frameStarts, frameEnds);
                    }
                }

                setRowNumbersForEmptyFrames(validRows);
                auto rowNumbersRange = melon::Range(rowNumbers_.data(), numRows);
                partition_->extractColumn(
                    valueIndex_, rowNumbersRange, resultOffset, result);
            }

        private:
            void setRowNumbersForEmptyFrames(const SelectivityVector &validRows) {
                if (validRows.isAllSelected()) {
                    return;
                }
                // Rows with empty (not-valid) frames have nullptr in the result.
                // So mark rowNumber to copy as kNullRow for it.
                invalidRows_.resizeFill(validRows.size(), true);
                invalidRows_.deselect(validRows);
                invalidRows_.applyToSelected([&](auto i) { rowNumbers_[i] = kNullRow; });
            }

            void setRowNumbersRespectNulls(
                const SelectivityVector &validRows,
                const BufferPtr &frameStarts,
                const BufferPtr &frameEnds) {
                if constexpr (TValue == ValueType::kFirst) {
                    auto rawFrameStarts = frameStarts->as<vector_size_t>();
                    validRows.applyToSelected(
                        [&](auto i) { rowNumbers_[i] = rawFrameStarts[i]; });
                } else {
                    auto rawFrameEnds = frameEnds->as<vector_size_t>();
                    validRows.applyToSelected(
                        [&](auto i) { rowNumbers_[i] = rawFrameEnds[i]; });
                }
            }

            void setRowNumbersIgnoreNulls(
                const SelectivityVector &validRows,
                const BufferPtr &frameStarts,
                const BufferPtr &frameEnds) {
                auto extractNullsResult = partition_->extractNulls(
                    valueIndex_, validRows, frameStarts, frameEnds, &nulls_);
                if (!extractNullsResult.has_value()) {
                    // There are no nulls in the column. Continue the processing with the
                    // function that respects nulls since it is more efficient.
                    return setRowNumbersRespectNulls(validRows, frameStarts, frameEnds);
                }

                auto leastFrame = extractNullsResult->first;
                auto frameSize = extractNullsResult->second;

                // first(last)Value functions return the first(last) non-null values for the
                // frame. Negate the bits in nulls_ so that nonNull bits are set instead.
                bits::negate(nulls_->asMutable<uint64_t>(), frameSize);
                auto rawNonNulls = nulls_->as<uint64_t>();

                auto rawFrameStarts = frameStarts->as<vector_size_t>();
                auto rawFrameEnds = frameEnds->as<vector_size_t>();
                validRows.applyToSelected([&](auto i) {
                    auto frameStart = rawFrameStarts[i];
                    auto frameEnd = rawFrameEnds[i];
                    // bits::findFirst(Last)Bit returns -1 if a set bit is not found.
                    // The function returns null for this case. -1 correctly maps to
                    // kNullRow as expected for rowNumbers_ extraction.
                    if constexpr (TValue == ValueType::kFirst) {
                        auto position = bits::findFirstBit(
                            rawNonNulls, frameStart - leastFrame, frameEnd - leastFrame + 1);
                        rowNumbers_[i] = (position == -1) ? -1 : position + leastFrame;
                    } else {
                        auto position = bits::findLastBit(
                            rawNonNulls, frameStart - leastFrame, frameEnd - leastFrame + 1);
                        rowNumbers_[i] = (position == -1) ? -1 : position + leastFrame;
                    }
                });
            }

            const bool ignoreNulls_;

            // Index of the first_value / last_value argument column in the input row
            // vector. This is used to retrieve column values from the partition data.
            column_index_t valueIndex_;

            const exec::WindowPartition *partition_;

            // The first_value, last_value functions directly write from the input column
            // to the resultVector using the extractColumn API specifying the rowNumber
            // mapping to copy between the 2 vectors. This variable is used for the
            // rowNumber vector across getOutput calls.
            std::vector<vector_size_t> rowNumbers_;

            // Used to extract nulls positions for the input value column if ignoreNulls
            // is set.
            BufferPtr nulls_;

            // Member variable re-used for setting null for empty frames.
            SelectivityVector invalidRows_;
        };
    } // namespace

    template<ValueType TValue>
    void registerFirstLastInternal(const std::string &name) {
        // T -> T
        std::vector<exec::FunctionSignaturePtr> signatures{
            exec::FunctionSignatureBuilder()
            .typeVariable("T")
            .returnType("T")
            .argumentType("T")
            .build(),
        };

        exec::registerWindowFunction(
            name,
            std::move(signatures),
            exec::WindowFunction::Metadata::defaultMetadata(),
            [](const std::vector<exec::WindowFunctionArg> &args,
               const TypePtr &resultType,
               bool ignoreNulls,
               pollux::memory::MemoryPool *pool,
               HashStringAllocator * /*stringAllocator*/,
               const pollux::core::QueryConfig & /*queryConfig*/)
        -> std::unique_ptr<exec::WindowFunction> {
                return std::make_unique<FirstLastValueFunction<TValue> >(
                    args, resultType, ignoreNulls, pool);
            });
    }

    void registerFirstValue(const std::string &name) {
        registerFirstLastInternal<ValueType::kFirst>(name);
    }

    void registerLastValue(const std::string &name) {
        registerFirstLastInternal<ValueType::kLast>(name);
    }
} // namespace kumo::pollux::window::prestosql
